/*
 * Copyright (c) 2018-2023 北京华恒盛世科技有限公司
 * QUICKPOOL is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/
#include "js.h"

#ifdef __cplusplus
extern "C" {
#endif

#define PREEMPTING    0x01

struct preemptive {
    JOBID_T jobid;
    HASH_TABLE_T *hash_prmable_job;
    unsigned int num_slot;
    int flags;
    time_t prm_time;
    unsigned int num_preempt_node;
    struct executeNode *preempt_nodes;
};

static HASH_TABLE_T *shashPreemptiveJob;

/**
 * @brief       policy: intialize preeption scheduling.
 *
 * @param[in]   queuelist    #1: preemption queue list
 *
 * @retval      0            #1: succeed
 *
 * @note        This function should be called after reading qp.queues.
 *              Make preemptive job hash table for scheduling.
 */
int
initialize_preemption(LIST_T *queuelist)
{
    if (NULL == shashPreemptiveJob) {
        shashPreemptiveJob = gf_hash_make(23);
    }

    return 0;
} // end function initialize_preemption

/**
 * @brief       policy: choose jobs for preemption.
 *
 * @param[in]   preemptqueue #1: preemptive queue
 * @param[out]  choosedjobs  #2: list of choosed preemtion jobs
 * @param[in]   maxprmjobs   #3: maximum number of preemption jobs
 *
 * @retval      >0           #1: succeed, number of preemption job
 * @retval      0            #2: failed, no job can be preempted
 *
 * @note        This function should be called during JS house keeping,
 *              to do job preemption from time to time.
 */
int
preempt_choose_jobs(struct queue *preemptqueue, LIST_T *choosedjobs, unsigned int maxprmjobs)
{
    if (0 == LIST_NUM_ENTS(preemptqueue->pending_jobs)) {

        log_message(DEBUG, PREEMPTION, "%s: no pending jobs in queue <%s> to do preempting.",
                    __func__, preemptqueue->queue_name);
        return 0;
    }

    LIST_T *listPreemptiveJob;
    listPreemptiveJob = gf_list_make("preemptive job list");

    char strJobId[32];

    struct job *preemptiveJob;
    for (preemptiveJob=(struct job *)preemptqueue->pending_jobs->back;
         preemptiveJob!=(struct job *)preemptqueue->pending_jobs;
         preemptiveJob=(struct job *)preemptiveJob->back) {

        if (!(preemptiveJob->job_status & JOB_STATUS_PEND)) {
            continue;
        }

        struct preemptive *prmJob;
        char key[32];
        snprintf(key, 32, "%lld", preemptiveJob->jobid);
        prmJob = (struct preemptive *)gf_hash_find(shashPreemptiveJob, key);
        if (NULL != prmJob && prmJob->flags & PREEMPTING) {
            continue;
        }

        if (0 == preemptiveJob->newReason) {

            // insert the preemptive job proxy into the end of list
            LIST_PROXY_T *pxy;
            pxy = (LIST_PROXY_T *)calloc(1, sizeof(struct list_proxy));
            if (NULL == pxy) {
                log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                            "calloc", sizeof(struct list_proxy));
                return 0;
            }

            pxy->subject = (void *)preemptiveJob;
            gf_list_insert(listPreemptiveJob, listPreemptiveJob, (LIST_T *)pxy);

            if (gf_get_logclass() & PREEMPTION) {
                log_message(DEBUG, PREEMPTION, "%s: job[%s] in queue <%s> can trigger preemption",
                            __func__, jobid_2string(preemptiveJob->jobid, strJobId),
                            preemptqueue->queue_name);
            }
        }
    }

    if (0 == LIST_NUM_ENTS(listPreemptiveJob)) {

        gf_list_free(listPreemptiveJob, NULL);

        log_message(DEBUG, PREEMPTION, "%s: No pending jobs to trigger preemption in queue <%s>.",
                    __func__, preemptqueue->queue_name);

        return 0;
    }

    LIST_ITERATOR_T iter;
    gf_listiterator_link(&iter, listPreemptiveJob);

    // Go through preemptive job and do preemption
    LIST_PROXY_T *pxy;
    for (pxy = (LIST_PROXY_T *)gf_listiterator_curentry(&iter);
         !gf_listiterator_isend(&iter);
         gf_listiterator_getnext(&iter, (LIST_T **)&pxy)) {

        preemptiveJob = (struct job *)pxy->subject;

        unsigned int numNeedSlot = preemptiveJob->shared->submit.num_job_slot;
        unsigned int num = 0;

        struct preemptive *prmJob;
        char strJobIdkey[32];
        snprintf(strJobIdkey, 32, "%lld", preemptiveJob->jobid);
        prmJob = (struct preemptive *)gf_hash_find(shashPreemptiveJob, strJobIdkey);
        if (NULL == prmJob) {

            prmJob = (struct preemptive *)calloc(1, sizeof(struct preemptive));
            if (NULL == prmJob) {
                log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                            "calloc", sizeof(struct preemptive));
                return 0;
            }

            prmJob->jobid = preemptiveJob->jobid;
            gf_hash_install(shashPreemptiveJob, strJobIdkey, (void *)prmJob);
        }

        struct queue *preemptableQueue;
        HASH_WALK_T walk;
        char *key;
        gf_hash_walk_start(preemptiveJob->job_queue->preemptable, &walk);
        while (NULL != (preemptableQueue = (struct queue *)gf_hash_walk(&walk, &key))) {

            if (0 == preemptableQueue->num_running_slot) {
                continue;
            }

            if (gf_get_logclass() & PREEMPTION) {
                log_message(DEBUG, PREEMPTION, "%s: trying to preempt starting jobs in queue <%s> for job [%s] in queue <%s>, need [%d] slots.",
	                    __func__, preemptableQueue->queue_name,
                            jobid_2string(preemptiveJob->jobid, strJobId),
                            preemptqueue->queue_name, numNeedSlot);
            }

            LIST_ITERATOR_T iter;
            gf_listiterator_link(&iter, glistStartedJob);

            struct job *preemptableJob;
            for (preemptableJob = (struct job *)gf_listiterator_curentry(&iter);
                 !gf_listiterator_isend(&iter);
                 gf_listiterator_getnext(&iter, (LIST_T **)&preemptableJob)) {

                if (preemptableJob->job_queue != preemptableQueue) {
                    continue;
                }

                if (NULL == preemptableJob->start_info) {
                    continue;
                }

                if (NULL == prmJob->hash_prmable_job) {
                    prmJob->hash_prmable_job = gf_hash_make(10);
                }

                snprintf(key, 32, "%lld", preemptableJob->jobid);
                struct job *tmpJob;
                tmpJob = (struct job *)gf_hash_find(prmJob->hash_prmable_job, key);
                if (NULL == tmpJob) {

                    num = num + preemptableJob->start_info->number_alloc_slot;

                    LIST_PROXY_T *pxy;
                    pxy = (LIST_PROXY_T *)calloc(1, sizeof(struct list_proxy));
                    if (NULL == pxy) {
                        log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                                    "calloc", sizeof(struct list_proxy));
                        return 0;
                    }

                    pxy->subject = (void *)preemptableJob;
                    gf_list_insert(choosedjobs, choosedjobs->forw, (LIST_T *)pxy);
                    gf_hash_install(prmJob->hash_prmable_job, key, (void *)preemptableJob);
                    prmJob->num_slot += preemptableJob->start_info->number_alloc_slot;

                    preemptableJob->preempting_job = preemptiveJob;

                    unsigned int i;
                    if (0 == prmJob->num_preempt_node) {
                        prmJob->num_preempt_node = preemptableJob->start_info->number_exec_node;
                        prmJob->preempt_nodes = (struct executeNode *)calloc(prmJob->num_preempt_node,
                                                                             sizeof(struct executeNode));
                        if (NULL == prmJob->preempt_nodes) {
                            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                                        "calloc", prmJob->num_preempt_node*sizeof(struct executeNode));
                            return 0;
                        }

                        for (i=0; i<prmJob->num_preempt_node; i++) {
                            prmJob->preempt_nodes[i].enode = preemptableJob->start_info->exec_nodes[i].enode;
                            prmJob->preempt_nodes[i].number_slot = preemptableJob->start_info->exec_nodes[i].number_slot;
                        }
                    } else {
                        unsigned int numNode = prmJob->num_preempt_node;
                        prmJob->num_preempt_node += preemptableJob->start_info->number_exec_node;
                        prmJob->preempt_nodes = (struct executeNode *)realloc(prmJob->preempt_nodes,
                                                                              prmJob->num_preempt_node*sizeof(struct executeNode));
                        if (NULL == prmJob->preempt_nodes) {
                            log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                                        "realloc", prmJob->num_preempt_node*sizeof(struct executeNode));
                            return 0;
                        }

                        for (i=numNode; i<prmJob->num_preempt_node; i++) {
                            prmJob->preempt_nodes[i].enode = preemptableJob->start_info->exec_nodes[i-numNode].enode;
                            prmJob->preempt_nodes[i].number_slot = preemptableJob->start_info->exec_nodes[i-numNode].number_slot;
                        }
                    }
                } else {

                    num = num + tmpJob->start_info->number_alloc_slot;

                    LIST_PROXY_T *pxy;
                    pxy = (LIST_PROXY_T *)calloc(1, sizeof(struct list_proxy));
                    if (NULL == pxy) {
                        log_message(ERR, NOCLASS, STR_FUNC_D_FAIL_M, __func__,
                                    "calloc", sizeof(struct list_proxy));
                        return 0;
                    }

                    pxy->subject = (void *)tmpJob;
                    gf_list_insert(choosedjobs, choosedjobs->forw, (LIST_T *)pxy);
                }

                if (gf_get_logclass() & PREEMPTION) {
                    log_message(DEBUG, PREEMPTION, "%s: job [%s] gives %d slots, got %d, need [%d] slots.",
                                __func__, jobid_2string(preemptableJob->jobid, strJobId),
                                preemptableJob->shared->submit.num_job_slot, num, numNeedSlot);
                }

                if (num >= numNeedSlot) {
                    break;
                }
            }

            if (num >= numNeedSlot) {
                break;
            }
        } // end while (preemptableQueue)
        gf_hash_walk_end(&walk);

        if (num >= numNeedSlot) {

            prmJob->prm_time = time(NULL);
            prmJob->flags |= PREEMPTING;

            break;
        } else { // Not enough slot for preemptive job, do not trigger preemption

            HASH_WALK_T walk;
            struct job *tmpJob;
            char *strJobKey;
            gf_hash_walk_start(prmJob->hash_prmable_job, &walk);
            while (NULL != (tmpJob = (struct job *)gf_hash_walk(&walk, &strJobKey))) {

                LIST_PROXY_T *pxy;
                pxy = (LIST_PROXY_T *)gf_list_search(choosedjobs, tmpJob, compare_proxy_job);
                if (NULL != pxy) {
                    gf_list_remove(choosedjobs, (LIST_T *)pxy);
                }

                tmpJob->preempting_job = NULL;
            }

            gf_hash_free(prmJob->hash_prmable_job, gf_free_none);
            prmJob->prm_time = 0;
            prmJob->flags = 0;
            prmJob->num_slot = 0;
        }

        if (LIST_NUM_ENTS(choosedjobs) >= maxprmjobs) {
            break;
        }
    } // end for (pxy)

    gf_list_free(listPreemptiveJob, NULL);

    return LIST_NUM_ENTS(choosedjobs);
} // end function preempt_choose_jobs

/**
 * @brief       get the number of preempting slots
 *
 * @param[in]   preemptingjob #1: preemptive queue
 * @param[in]   usablenode    #2: maximum number of preemption jobs
 *
 * @retval      >0           #1: succeed, number of preempting slots
 * @retval      0            #2: failed, job is not preemption slots
 *
 * @note        This function should be called during scheduling,
 *              to get the number of preemptive slots of a preemptive job.
 */
int
get_preemption_slot(struct job *preemptingjob, struct node *usablenode)
{
    struct preemptive *prmJob;
    char key[32];
    snprintf(key, 32, "%lld", preemptingjob->jobid);
    prmJob = (struct preemptive *)gf_hash_find(shashPreemptiveJob, key);
    if (NULL == prmJob || 0 == prmJob->num_preempt_node) {
        // job is not preempting slots from other jobs
        return 0;
    }

    int numPreemptSlot=0;
    unsigned int i;
    for (i=0; i<prmJob->num_preempt_node; i++) {
        if (prmJob->preempt_nodes[i].enode != usablenode) {
            continue;
        }

        numPreemptSlot += prmJob->preempt_nodes[i].number_slot;
    }

    return numPreemptSlot;
} // end function get_preemption_slot

#ifdef __cplusplus
}
#endif
